import re
import uuid

import pandas as pd
from tqdm import tqdm

from tree_rag.adapters import embedding
from tree_rag.dataclasses.api import Clarify, Knowledge
from tree_rag.dataclasses.knowledge import Graph, Scene
from tree_rag.es import index_clarify, index_document


async def prapare_data_from_excel(
    sheet: str, dataset_id: str, tenant_id: str, channel_id: str, ktype: str = "normal"
):
    def parse(tags) -> Graph:
        tags = [
            line.strip()
            for line in tags.split("\n")
            if line.strip() and isinstance(line, str)
        ]
        level_1 = []
        level_2 = []
        scene = []
        keywords = []

        for tag in tags:
            if re.match(r"一级事项引导[：:]", tag):
                cur_level_1 = re.split(r"[,，]", re.sub(r"一级事项引导[：:]", "", tag))
                level_1 = [item.strip() for item in cur_level_1 if item]

            if re.match(r"二级事项引导[：:]", tag):
                level_2 = re.split(r"[,，]", re.sub(r"二级事项引导[：:]", "", tag))

            if re.match(r"场景引导[：:]", tag):
                scene = re.split(r"[,，]", re.sub(r"场景引导[：:]", "", tag))
                scene = [
                    Scene(label=item, category=scene2category[item]).model_dump()
                    for item in scene
                    if item in scene2category
                ]

            if re.match(r"关键词[：:]", tag):
                keywords = re.split(r"[,，、]", re.sub(r"关键词[：:]", "", tag))

        return Graph(level_1=level_1, level_2=level_2, scene=scene, keywords=keywords)

    knowledge_sheet = pd.read_excel(sheet, sheet_name="QA")
    try:
        # 没有场景引导表
        category2clarify = {}
        scene2category = {}

        clarify_sheet = pd.read_excel(sheet, sheet_name="场景引导")
        for _, row in clarify_sheet.iterrows():
            for item in row["场景引导"].split("\n"):
                scene2category[item] = row["场景引导类型"]
            category2clarify[row["场景引导类型"]] = row["澄清话术"]
    except ValueError:
        # 没有场景引导表
        category2clarify = {}
        scene2category = {}

    for _, row in tqdm(knowledge_sheet.iterrows(), desc="Uploading knowledge"):
        if "tag" not in row and "标签" not in row:
            level_1_str = row["一级引导"] if "一级引导" in row else row["level_1"]
            level_2_str = row["二级引导"] if "二级引导" in row else row["level_2"]
            scene = row["场景"] if "场景" in row else row["scene"]
            scene = Scene.fix_str(scene, scene2category)
            keywords = row["关键词"] if "关键词" in row else row["keywords"]

            graph = Graph(
                level_1=level_1_str,
                level_2=level_2_str,
                scene=scene,
                keywords=keywords,
            )

        else:
            tags = row["tag"] if "tag" in row else row["标签"]
            if not isinstance(tags, str):
                continue
            graph = parse(tags)

        if "id" not in row:
            row["id"] = str(uuid.uuid4())

        else:
            try:
                str(uuid.UUID(row["id"]))
            except:
                row["id"] = str(row["id"])
                row["id"] = row["id"] + "0" * (32 - len(row["id"]))

        question = row["question"] if "question" in row else row["问题"]
        if "answer" not in row and "答案" not in row:
            answer = ""
        else:
            answer = row["answer"] if "answer" in row else row["答案"]
        knowledge = Knowledge(
            dataset_id=dataset_id,
            segment_id=uuid.uuid4().hex,
            desheng_id=row["id"],
            tenant_id=tenant_id,
            channel_id=channel_id,
            question=question,
            answer=answer,
            graph=graph,
            ktype=ktype,
        )
        question_embedding = await embedding(knowledge.question)
        answer_embedding = await embedding(knowledge.answer)
        await index_document(knowledge, question_embedding, answer_embedding)

    for category, words in tqdm(category2clarify.items(), desc="Uploading clarify"):
        if pd.isna(words):
            continue
        clarify = Clarify(
            tenant_id=tenant_id,
            channel_id=[channel_id],
            category=category,
            words=words,
        )
        await index_clarify(clarify)